import os


class YoloLableFormat(object):
    suffix = ".txt"

    def __init__(self) -> None:
        self.shapes = []
        self.categories = "class.yaml"
        self.label_map = dict()
        self.label_set = set()

    def __xywh2xyxy(self, split_data, h, w):
        """中心坐标宽高转左上右下坐标"""
        cx, cy, bw, bh = split_data
        cx = float(cx) * w
        cy = float(cy) * h
        bw = float(bw) * w
        bh = float(bh) * h
        x1 = cx - bw / 2
        y1 = cy - bh / 2
        x2 = cx + bw / 2
        y2 = cy + bh / 2
        return [[x1, y1], [x2, y2]]

    def __string2class(self, label_str):
        """字符串转类别"""
        for key, value in self.label_map["names"].items():
            if value == label_str:
                return key
        return -1

    def get_label_color(self, label_str):
        color = None
        cls = self.__string2class(label_str)
        if cls in self.label_map["colors"]:
            color = self.label_map["colors"][cls]
        return eval(color)

    def read(self, filename, image_size):
        """加载标签"""
        h, w = image_size
        ret = []
        if os.path.exists(filename):
            with open(filename) as obj:
                data = obj.read().splitlines()
                for oneline in data:
                    split_data = oneline.split(" ")
                    cls = int(split_data[0])
                    color = None
                    cls_name = ""
                    if not self.label_map:
                        cls_name = split_data[0]
                    else:
                        if "colors" in self.label_map and cls in self.label_map["colors"]:
                            color = self.label_map["colors"][cls]
                        cls_name =  self.label_map["names"][cls]
                    ret.append({
                        "label": cls_name,
                        "color": color,
                        "points": self.__xywh2xyxy(split_data[1:], h, w),
                        "group_id": [None],
                        "attributes": [],
                        "shape_type": 2
                    })
        self.shapes = ret

    def write(self, filename, shapes, image_w, image_h):
        """保存标签数据"""
        with open(filename, "w") as obj:
            for one_shape in shapes:
                label = self.__string2class(one_shape["label"])
                (x1, y1), (x2, y2) = one_shape["points"]
                cx = (x1 + x2) / 2 / image_w
                cy = (y1 + y2) / 2 / image_h
                w = abs(x1 - x2) / image_w
                h = abs(y1 - y2) / image_h
                write_str = " ".join([str(label), str(cx), str(cy), str(w), str(h)]) + "\n"
                obj.write(write_str)

    def is_label_file(self, filename):
        """是否是标签文件"""
        return os.path.splitext(filename)[0].lower() == YoloLableFormat.suffix

    def writeCategoriesFile(self, save_path, save_info):
        """保存类别信息"""
        categories_file = os.path.join(save_path, self.categories)
        for onelabel in save_info:
            if "name" in onelabel and onelabel["name"] in self.label_set:
                continue
            else:
                self.label_set.add(onelabel["name"])
                cls = len(self.label_map["names"])
                if "color" in onelabel:
                    self.label_map["colors"][cls] = str(onelabel["color"])
                self.label_map["names"][cls] = onelabel["name"]
        with open(categories_file, "w") as obj:
            import yaml
            yaml.dump(self.label_map, obj)

    def readCategoriesFile(self, file_path):
        """读取类别信息"""
        categories_file = os.path.join(file_path, self.categories)
        data = None
        if os.path.exists(categories_file):
            with open(categories_file) as obj:
                import yaml
                self.label_map = yaml.safe_load(obj)
                data = []
                for key, value in self.label_map["names"].items():
                    self.label_set.add(value)
                    color = None
                    if key in self.label_map["colors"]:
                        color = eval(self.label_map["colors"][key])
                    data.append({
                        "name": value,
                        "color": color,
                        "attributes": []
                    })
        return data
